#!/usr/bin/python

import numpy as np
import pylab as p
import math as m
import os
import sys

from VISUALIZATION__DEFINE import *
from SIMULATION__DEFINE import *

v = visualization()
if not os.path.isdir(os.path.expanduser(v.plot_folder)):
	os.mkdir(os.path.expanduser(v.plot_folder))


try:
	os.chdir(os.path.expanduser(v.run))
	print "Starting "+v.run+" visualization"
except:
	print "run folder not defined correctly"
	sys.exit(1)

s = simulation('.')

# plot timeseries
data = s.getTimeseries()
for j in v.ts_fields:
	filename = os.path.expanduser(v.plot_folder)+"ts_"+j+".png"
	print "\tplotting "+os.path.basename(filename)

	p.figure()
	p.plot(data['time'],data[j], linewidth=2)
	p.xlabel('Non-dimensional time', weight='bold')
	p.ylabel('Non-dimensional '+j, weight='bold')
	p.savefig(filename)
	p.close()

for i in range(len(s.iterations)):
	iter = s.iterations[i]
	time = s.times[i]
	print "in iter "+str(iter)

	# plot profiles
	first = 1
	for j in v.profiles:
		filename = os.path.expanduser(v.plot_folder)+j+str(iter)+".png"
		if not os.path.isfile(filename):
			print "\tplotting "+os.path.basename(filename)
			if first:
				print "loading data"
				data = s.getProfiles(iter)
				first = 0
			p.figure()
			p.plot(data['depth'], data[j], linewidth=2)
			p.xlabel('Non-dimensional radius', weight='bold')
			p.ylabel('Non-dimensional '+j, weight='bold')
			p.title('t = '+time, weight='bold')
			p.savefig(filename)
			p.close()

	# plot grids
	first = 1
	for j in v.fields:
		filename = os.path.expanduser(v.plot_folder)+"field_"+j+str(iter)+".png"
		if not os.path.isfile(filename):
			print "\tplotting "+os.path.basename(filename)
			if first:
				print "loading data"
				data = s.getGrids(iter)
				first = 0

			D = data[j].reshape(s.nx,s.nz)
			D = D.T
			D = D[::-1,:]

			N = np.sqrt(np.power(data['vx'],2)+np.power(data['vz'],2))
			VX = data['vx']
			VZ = data['vz']
			VX[np.nonzero(VX+VZ != 0)] /= N[np.nonzero(VX+VZ != 0)]
			VZ[np.nonzero(VX+VZ != 0)] /= N[np.nonzero(VX+VZ != 0)]
	
			p.figure(figsize=(s.aspect_ratio*2.5,2.5))
			p.title('t = '+time, weight='bold')
			IM = p.pcolor(data['x'],data['z'],D)
			CB = p.colorbar(IM,)
			p.axis('scaled')
			CB.ax.set_ylabel(j, weight='bold')
			p.quiver(data['x'][::10],data['z'][::10],VX[::10,::10],VZ[::10,::10], pivot='middle')
			p.savefig(filename)
			p.close()
